{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SUknhuHKyc-E"
   },
   "source": [
    "# <center>OpenAI agent pattern: evaluator optimizer agent</center>\n",
    "\n",
    "A starter guide for building an agent which iteratively generates an output based on LLM feedback using the `openai-agents` library.\n",
    "\n",
    "When creating LLM outputs, often times the first generation is unsatisfactory. You can use an agentic loop to iteratively improve the output by asking an LLM to give feedback, and then use the feedback to improve the output.\n",
    "\n",
    "In the following example, we'll build a financial report system using this pattern:\n",
    "1.  **Report Agent (Generation):** Creates a report on a particular stock ticker.\n",
    "2.  **Evaluator Agent (Feedback):** Evaluates the report and provides feedback on what to improve."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n69HR7eJswNt"
   },
   "source": [
    "### Install Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install base libraries for OpenAI\n",
    "!pip install -q openai openai-agents pydantic\n",
    "\n",
    "# Install optional libraries for OpenInference/OpenTelemetry tracing\n",
    "!pip install -q arize-phoenix-otel openinference-instrumentation-openai-agents openinference-instrumentation-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jQnyEnJisyn3"
   },
   "source": [
    "### Setup Keys\n",
    "\n",
    "Add your OpenAI API key to the environment variable `OPENAI_API_KEY`.\n",
    "\n",
    "Copy your Phoenix `API_KEY` from your settings page at [app.phoenix.arize.com](https://app.phoenix.arize.com)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass(\"🔑 Enter your OpenAI API key: \")\n",
    "\n",
    "if \"PHOENIX_API_KEY\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_API_KEY\"] = getpass(\"🔑 Enter your Phoenix API key: \")\n",
    "\n",
    "if \"PHOENIX_COLLECTOR_ENDPOINT\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_COLLECTOR_ENDPOINT\"] = getpass(\"🔑 Enter your Phoenix Collector Endpoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kfid5cE99yN5"
   },
   "source": [
    "### Setup Tracing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from phoenix.otel import register\n",
    "\n",
    "tracer_provider = register(\n",
    "    project_name=\"openai-agents\",\n",
    "    auto_instrument=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the agent\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from textwrap import dedent\n",
    "from typing import Literal\n",
    "\n",
    "from agents import Agent, Runner, TResponseInputItem\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "CATALYSTS = \"\"\"topline revenue growth, margin expansion, moat expansion, free cash flow generation, usage, pricing, distribution, share buyback, dividend, new products, regulation, competition, management team, mergers, acquisitions, analyst ratings, trading volume, technical indicators, price momentum\"\"\"\n",
    "\n",
    "\n",
    "class EvaluationFeedback(BaseModel):\n",
    "    feedback: str = Field(\n",
    "        description=f\"What is missing from the research report on positive and negative catalysts for a particular stock ticker. Catalysts include changes in {CATALYSTS}.\"\n",
    "    )\n",
    "    score: Literal[\"pass\", \"needs_improvement\", \"fail\"] = Field(\n",
    "        description=\"A score on the research report. Pass if the report is complete and contains at least 3 positive and 3 negative catalysts for the right stock ticker, needs_improvement if the report is missing some information, and fail if the report is completely wrong.\"\n",
    "    )\n",
    "\n",
    "\n",
    "report_agent = Agent(\n",
    "    name=\"Catalyst Report Agent\",\n",
    "    instructions=dedent(\n",
    "        \"\"\"You are a research assistant specializing in stock research. Given a stock ticker, generate a report of 3 positive and 3 negative catalysts that could move the stock price in the future in 50 words or less.\"\"\"\n",
    "    ),\n",
    "    model=\"gpt-4.1\",\n",
    ")\n",
    "\n",
    "evaluation_agent = Agent(\n",
    "    name=\"Evaluation Agent\",\n",
    "    instructions=dedent(\n",
    "        \"\"\"You are a senior financial analyst. You will be provided with a stock research report with positive and negative catalysts. Your task is to evaluate the report and provide feedback on what to improve.\"\"\"\n",
    "    ),\n",
    "    model=\"gpt-4.1\",\n",
    "    output_type=EvaluationFeedback,\n",
    ")\n",
    "\n",
    "report_feedback = \"fail\"\n",
    "input_items: list[TResponseInputItem] = [{\"content\": \"AAPL\", \"role\": \"user\"}]\n",
    "\n",
    "while report_feedback != \"pass\":\n",
    "    report = await Runner.run(report_agent, input_items)\n",
    "    print(\"### REPORT ###\")\n",
    "    print(report.final_output)\n",
    "    input_items = report.to_input_list()\n",
    "\n",
    "    evaluation = await Runner.run(evaluation_agent, str(report.final_output))\n",
    "    evaluation_feedback = evaluation.final_output_as(EvaluationFeedback)\n",
    "    print(\"### EVALUATION ###\")\n",
    "    print(str(evaluation_feedback))\n",
    "    report_feedback = evaluation_feedback.score\n",
    "\n",
    "    if report_feedback != \"pass\":\n",
    "        print(\"Re-running with feedback\")\n",
    "        input_items.append({\"content\": f\"Feedback: {evaluation_feedback.feedback}\", \"role\": \"user\"})"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
